import csv
import hashlib
import os
from pathlib import Path
import pandas as pd
from tqdm import tqdm # 用于显示进度条，需要 pip install tqdm
from typing import Union # <--- 添加这一行导入

# ↓↓↓ 修改这一行的类型提示 ↓↓↓
def calculate_sha256(filepath: Path) -> Union[str, None]:
    """计算文件的 SHA-256 哈希值"""
    try:
        hasher = hashlib.sha256()
        with open(filepath, 'rb') as file:
            while True:
                # 分块读取以处理大文件
                chunk = file.read(8192)
                if not chunk:
                    break
                hasher.update(chunk)
        return hasher.hexdigest()
    except FileNotFoundError:
        print(f"错误：文件未找到 {filepath}")
        return None
    except Exception as e:
        print(f"计算哈希时出错 {filepath}: {e}")
        return None

# --- 下面的代码保持不变 ---

def index_dataset(dataset_dir: Path, image_extensions: set) -> dict[str, list[str]]:
    """
    遍历数据集目录，计算图片文件的哈希值并建立索引。
    返回一个字典，键是哈希值，值是包含该哈希值的文件路径列表。
    """
    hash_to_paths = {}
    print(f"正在索引数据集目录: {dataset_dir} ...")

    # 使用 rglob 递归查找所有文件
    all_files = list(dataset_dir.rglob('*'))

    # 使用 tqdm 显示进度
    for item in tqdm(all_files, desc="索引数据集文件"):
        try:
            # 确保是文件并且扩展名在允许的列表中
            if item.is_file() and item.suffix.lower() in image_extensions:
                file_hash = calculate_sha256(item)
                if file_hash:
                    # 如果哈希值已存在，将路径添加到列表中；否则创建新列表
                    hash_to_paths.setdefault(file_hash, []).append(str(item))
        except Exception as e:
            print(f"处理文件 {item} 时跳过，出现错误: {e}")
            continue # 继续处理下一个文件

    print(f"数据集索引完成，共找到 {len(hash_to_paths)} 个独特的图片哈希值。")
    return hash_to_paths

if __name__ == "__main__":
    # --- 配置区域 ---
    # !!! 请将下面的路径修改为你实际的目录 !!!
    QUERY_IMAGE_DIR = Path(r"/home/zhangyichi/dataset/OIA-ODIR/validation/image")  # 你需要检查的图片所在的目录
    DATASET_DIR = Path(r"/home/zhangyichi/dataset/OIA-ODIR/Training Set/Images")      # 庞大的原始数据集根目录

    # 定义要处理的图片文件扩展名（小写）
    IMAGE_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.bmp', '.gif', '.tif', '.tiff'}
    # --- 配置区域结束 ---

    # 检查路径是否存在
    if not QUERY_IMAGE_DIR.is_dir():
        print(f"错误：查询图片目录不存在: {QUERY_IMAGE_DIR}")
        exit()
    if not DATASET_DIR.is_dir():
        print(f"错误：数据集目录不存在: {DATASET_DIR}")
        exit()

    # 1. 对大数据集建立哈希索引
    dataset_index = index_dataset(DATASET_DIR, IMAGE_EXTENSIONS)

    if not dataset_index:
        print("未能从数据集中索引任何图片，程序退出。")
        exit()

    # 2. 遍历查询目录中的图片并进行比较
    found_matches = []
    print(f"\n正在检查查询目录中的图片: {QUERY_IMAGE_DIR} ...")

    query_files = list(QUERY_IMAGE_DIR.rglob('*')) # 同样递归查找

    for query_item in tqdm(query_files, desc="检查查询图片"):
         try:
            # 确保是文件并且扩展名在允许的列表中
            if query_item.is_file() and query_item.suffix.lower() in IMAGE_EXTENSIONS:
                query_hash = calculate_sha256(query_item)
                if query_hash:
                    # 在数据集中查找哈希值
                    if query_hash in dataset_index:
                        original_paths = dataset_index[query_hash]
                        found_matches.append({
                            "query_image": str(query_item),
                            "original_paths": original_paths
                        })
                    # else:
                        # print(f"未找到匹配: {query_item.name}") # 如果需要，可以取消注释
         except Exception as e:
            print(f"处理查询文件 {query_item} 时跳过，出现错误: {e}")
            continue # 继续处理下一个文件


    # 3. 报告结果
    print("\n--- 匹配结果 ---")
    output_csv = r'/home/zhangyichi/dataset/OIA-ODIR/answer.csv'
    df = pd.read_csv(r'/home/zhangyichi/dataset/OIA-ODIR/Training Set/Annotation/train single.csv')
    print(df[df['Fundus'] == '0_left.jpg'])
    if found_matches:
        print(f"总共找到 {len(found_matches)} 个匹配项：")
        with open(output_csv,mode='w', newline='', encoding='utf-8') as f:
            writer = csv.writer(f)
            header = ['ID', 'raw_name', 'train_name',
             'N', 'D', 'G', 'C', 'A', 'H', 'M', 'O']
            writer.writerow(header)
            id_counter = 0
            for match in found_matches:
                raw_name = match['query_image'].split('/')[-1]
                train_name = match['original_paths'][0].split('/')[-1]
                # print(raw_name)
                match_column = df.loc[df['Fundus'] == train_name,['N', 'D', 'G', 'C', 'A', 'H', 'M', 'O']].iloc[0]
                row = [
                    id_counter,
                    raw_name,
                    train_name,
                    match_column['N'],
                    match_column['D'],
                    match_column['G'],
                    match_column['C'],
                    match_column['A'],
                    match_column['H'],
                    match_column['M'],
                    match_column['O']
                ]
                writer.writerow(row)
                id_counter += 1

    else:
        print("在数据集中没有找到任何与查询图片完全匹配的文件。")

    print("\n检查完成。")